import os
import time
import tkinter as tk
from tkinter import messagebox, filedialog, colorchooser
import pyautogui
from pynput import keyboard
from PIL import Image, ImageDraw, ImageFont
import threading

SAVE_DIR = "screenshots"
FONT_PATH = "simsun.ttc"
FONT_SIZE = 36
TEXT_COLOR = (255, 255, 255, 64)
LISTENER = None

def take_screenshot():

    try:
        watermark_text = watermark_entry.get()

        timestamp = time.strftime("%Y%m%d_%H%M%S")
        screenshot_path = os.path.join(SAVE_DIR, f"screenshot_{timestamp}.png")

        if not os.path.exists(SAVE_DIR):
            os.makedirs(SAVE_DIR)

        screenshot = pyautogui.screenshot()
        screenshot.save(screenshot_path)
        print(f"截图已保存到 {screenshot_path}")

        if add_watermark_var.get():
            add_watermark(screenshot_path, watermark_text, FONT_PATH, FONT_SIZE, TEXT_COLOR,ANGLE)

        threading.Thread(target=show_non_blocking_messagebox, args=("成功", f"截图已保存: {screenshot_path}")).start()
    except Exception as e:
        threading.Thread(target=show_non_blocking_messagebox, args=("错误", f"截图时出错: {e}")).start()

def add_watermark(image_path, watermark_text, font_path, font_size, text_color,angle=30):

    try:
        image = Image.open(image_path).convert("RGBA")
        
        watermark = Image.new("RGBA", image.size, (0, 0, 0, 0))
        
        draw = ImageDraw.Draw(watermark)
        
        try:
            font = ImageFont.truetype(font_path, font_size)
        except IOError:
            print(f"字体文件 {font_path} 未找到，使用默认字体。")
            font = ImageFont.load_default()


        bbox = font.getbbox(watermark_text)
        text_width = bbox[2] - bbox[0]
        text_height = bbox[3] - bbox[1]



        text_layer = Image.new("RGBA", (text_width, text_height), (0, 0, 0, 0))
        text_draw = ImageDraw.Draw(text_layer)
        text_draw.text((0, 0), watermark_text, font=font, fill=text_color)
        rotated_layer = text_layer.rotate(
            angle, 
            expand=True, 
            fillcolor=(0, 0, 0, 0),  
            resample=Image.BILINEAR
        )
        # rotated_layer.show()
        rw, rh = rotated_layer.size 
        step_x = int(rw * 1.5)
        step_y = int(rh * 1.5)
        for x in range(-rw, image.width + rw, step_x):
            for y in range(-rh, image.height + rh, step_y):
                
                pos = (
                    x - rw // 2,
                    y - rh // 2
                )
                watermark.paste(rotated_layer, pos)
                
        watermarked = Image.alpha_composite(image, watermark)
        if image_path.lower().endswith(('.jpg', '.jpeg')):
            watermarked = watermarked.convert('RGB')
        watermarked.show()
        watermarked.save(image_path)
        print(f"水印已添加到 {image_path}")
    except Exception as e:
        threading.Thread(target=show_non_blocking_messagebox, args=("错误", f"添加水印时出错: {e}")).start()

def show_non_blocking_messagebox(title, message):

    root = tk.Tk()
    root.withdraw()         
    root.attributes("-topmost", True)
    messagebox.showinfo(title, message)
    root.destroy()

def on_f2_pressed():

    take_screenshot()

def on_press(key):

    try:
        if key == keyboard.Key.f2:
            on_f2_pressed()
    except AttributeError:
        pass

def start_keyboard_listener():

    global LISTENER
    LISTENER = keyboard.Listener(on_press=on_press)
    LISTENER.start()

def select_save_dir():

    global SAVE_DIR
    save_dir = filedialog.askdirectory(title="选择截图保存目录")
    if save_dir:
        SAVE_DIR = save_dir
        save_dir_label.config(text=f"保存目录: {SAVE_DIR}")

def select_color():

    global TEXT_COLOR
    color = colorchooser.askcolor(title="选择水印颜色")
    if color[1]:
        r, g, b = map(int, color[0])
        TEXT_COLOR = (r, g, b, 64)
        color_label.config(text=f"水印颜色: {color[1]}")
def update_angle(val):
    global ANGLE
    ANGLE = int(float(val))
    angle_value_label.config(text=f"当前角度: {ANGLE}°")

root = tk.Tk()
root.title("后台按键截图工具")
root.geometry("400x400")


tk.Label(root, text="按下 F2 键截图（支持后台）").pack(pady=10)


tk.Label(root, text="水印文字:").pack(pady=5)
watermark_entry = tk.Entry(root, width=30)
watermark_entry.insert(0, "202472610119 于洪朋")  # 默认水印文字
watermark_entry.pack(pady=5)


save_dir_label = tk.Label(root, text=f"保存目录: {SAVE_DIR}")
save_dir_label.pack(pady=5)

select_dir_button = tk.Button(root, text="选择保存目录", command=select_save_dir)
select_dir_button.pack(pady=5)

add_watermark_var = tk.BooleanVar(value=True)
add_watermark_check = tk.Checkbutton(root, text="添加水印", variable=add_watermark_var)
add_watermark_check.pack(pady=5)

color_label = tk.Label(root, text="水印颜色: 默认白色")
color_label.pack(pady=5)

select_color_button = tk.Button(root, text="选择水印颜色", command=select_color)
select_color_button.pack(pady=5)
angle_frame = tk.Frame(root)
angle_frame.pack(pady=10)
tk.Label(angle_frame, text="水印旋转角度:").pack(side=tk.LEFT)
angle_slider = tk.Scale(
    angle_frame,
    from_=-180,
    to=180,
    orient=tk.HORIZONTAL,
    length=200,
    showvalue=0, 
    command=update_angle
)
angle_slider.set(45)  
angle_slider.pack(side=tk.LEFT, padx=5)
angle_value_label = tk.Label(angle_frame, text="当前角度: 0°")
angle_value_label.pack(side=tk.LEFT)
start_keyboard_listener()

root.mainloop()